#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
标注工具栏组件 - 提供各种标注工具
"""

from PyQt5.QtWidgets import (
    QToolBar, QAction, QActionGroup, QPushButton, QComboBox,
    QLabel, QColorDialog, QWidget, QHBoxLayout, QVBoxLayout, 
    QButtonGroup, QToolButton, QMenu, QApplication
)
from PyQt5.QtCore import Qt, QSize, pyqtSignal
from PyQt5.QtGui import QIcon, QColor, QPixmap, QPainter, QPen

class ColorButton(QPushButton):
    """颜色选择按钮"""
    
    colorChanged = pyqtSignal(QColor)
    
    def __init__(self, color=Qt.red, parent=None):
        super().__init__(parent)
        self.setColor(color)
        self.clicked.connect(self.onButtonClicked)
        self.setFixedSize(24, 24)
    
    def setColor(self, color):
        """设置颜色"""
        self.color = color
        
        # 创建一个小图像用作图标
        pixmap = QPixmap(16, 16)
        pixmap.fill(self.color)
        
        self.setIcon(QIcon(pixmap))
        self.update()
    
    def onButtonClicked(self):
        """按钮点击事件"""
        color = QColorDialog.getColor(self.color, self, "选择标注颜色")
        if color.isValid():
            self.setColor(color)
            self.colorChanged.emit(color)


class AnnotationToolbar(QToolBar):
    """标注工具栏"""
    
    # 自定义信号
    tool_changed = pyqtSignal(str)  # 工具变更
    color_changed = pyqtSignal(QColor)  # 颜色变更
    class_changed = pyqtSignal(str)  # 标签类别变更
    
    def __init__(self, parent=None):
        """初始化"""
        super().__init__("标注工具栏", parent)
        
        # 设置工具栏属性
        self.setMovable(False)
        self.setIconSize(QSize(32, 32))
        self.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
        
        # 存储工具动作
        self.tool_actions = {}
        
        # 初始化默认值
        self.current_tool = "select"
        self.current_color = QColor(Qt.red)
        self.current_class = ""
        
        # 初始化标注工具
        self._init_tools()
        
        # 添加颜色选择器
        self.add_color_selector()
        
        # 添加类别选择器
        self.add_class_selector()
    
    def _init_tools(self):
        """初始化标注工具"""
        # 选择工具
        select_action = QAction(QIcon("icons/select.png"), "选择", self)
        select_action.setCheckable(True)
        select_action.setChecked(True)  # 默认选中
        select_action.setData("select")
        self.addAction(select_action)
        self.tool_actions["select"] = select_action
        
        # 添加分隔符
        self.addSeparator()
        
        # 矩形工具
        rect_action = QAction(QIcon("icons/rectangle.png"), "矩形", self)
        rect_action.setCheckable(True)
        rect_action.setData("rectangle")
        self.addAction(rect_action)
        self.tool_actions["rectangle"] = rect_action
        
        # 多边形工具
        polygon_action = QAction(QIcon("icons/polygon.png"), "多边形", self)
        polygon_action.setCheckable(True)
        polygon_action.setData("polygon")
        self.addAction(polygon_action)
        self.tool_actions["polygon"] = polygon_action
        
        # 旋转矩形工具
        rot_rect_action = QAction(QIcon("icons/rotated_rectangle.png"), "旋转矩形", self)
        rot_rect_action.setCheckable(True)
        rot_rect_action.setData("rotated_rectangle")
        self.addAction(rot_rect_action)
        self.tool_actions["rotated_rectangle"] = rot_rect_action
        
        # 点标注工具
        point_action = QAction(QIcon("icons/point.png"), "点标注", self)
        point_action.setCheckable(True)
        point_action.setData("point")
        self.addAction(point_action)
        self.tool_actions["point"] = point_action
        
        # 添加分隔符
        self.addSeparator()
        
        # 智能分割工具
        smart_segment_action = QAction(QIcon("icons/smart_segment.png"), "智能分割", self)
        smart_segment_action.setCheckable(True)
        smart_segment_action.setData("smart_segment")
        self.addAction(smart_segment_action)
        self.tool_actions["smart_segment"] = smart_segment_action
        
        # 创建工具组
        tool_group = QActionGroup(self)
        for action in self.tool_actions.values():
            tool_group.addAction(action)
        
        # 连接信号
        tool_group.triggered.connect(self._on_tool_triggered)
    
    def add_color_selector(self):
        """添加颜色选择器"""
        # 创建颜色按钮的容器
        color_widget = QWidget()
        color_layout = QVBoxLayout(color_widget)
        color_layout.setContentsMargins(0, 0, 0, 0)
        
        # 添加标签
        color_label = QLabel("颜色")
        color_layout.addWidget(color_label, alignment=Qt.AlignCenter)
        
        # 创建颜色按钮
        self.color_button = ColorButton(self.current_color)
        self.color_button.colorChanged.connect(self._on_color_changed)
        
        color_layout.addWidget(self.color_button, alignment=Qt.AlignCenter)
        
        # 添加到工具栏
        self.addWidget(color_widget)
    
    def add_class_selector(self):
        """添加类别选择器"""
        # 创建类别选择器的容器
        class_widget = QWidget()
        class_layout = QVBoxLayout(class_widget)
        class_layout.setContentsMargins(0, 0, 0, 0)
        
        # 添加标签
        class_label = QLabel("标签类别")
        class_layout.addWidget(class_label, alignment=Qt.AlignCenter)
        
        # 创建下拉框
        self.class_combo = QComboBox()
        self.class_combo.setFixedWidth(100)  # 设置固定宽度
        self.class_combo.currentTextChanged.connect(self._on_class_changed)
        
        class_layout.addWidget(self.class_combo, alignment=Qt.AlignCenter)
        
        # 创建新增类别按钮
        self.add_class_button = QPushButton("添加类别")
        self.add_class_button.clicked.connect(self._on_add_class_clicked)
        
        class_layout.addWidget(self.add_class_button, alignment=Qt.AlignCenter)
        
        # 添加到工具栏
        self.addWidget(class_widget)
    
    def _on_tool_triggered(self, action):
        """工具选择事件"""
        # 获取选中的工具
        self.current_tool = action.data()
        
        # 发射信号
        self.tool_changed.emit(self.current_tool)
    
    def _on_color_changed(self, color):
        """颜色变更事件"""
        self.current_color = color
        
        # 发射信号
        self.color_changed.emit(color)
    
    def _on_class_changed(self, text):
        """类别变更事件"""
        self.current_class = text
        
        # 发射信号
        self.class_changed.emit(text)
    
    def _on_add_class_clicked(self):
        """添加类别按钮点击事件"""
        # 这个事件需要在主窗口中处理
        # 发射信号通知主窗口显示添加类别对话框
        pass
    
    def set_classes(self, classes):
        """
        设置可用的类别列表
        
        Args:
            classes (list): 类别列表
        """
        # 保存当前选择
        current_text = self.class_combo.currentText()
        
        # 清除当前选项
        self.class_combo.clear()
        
        # 添加类别
        for cls in classes:
            if isinstance(cls, dict) and "name" in cls:
                self.class_combo.addItem(cls["name"])
            elif isinstance(cls, str):
                self.class_combo.addItem(cls)
        
        # 恢复选择
        index = self.class_combo.findText(current_text)
        if index >= 0:
            self.class_combo.setCurrentIndex(index)
        elif self.class_combo.count() > 0:
            self.class_combo.setCurrentIndex(0)
    
    def get_current_tool(self):
        """获取当前选中的工具"""
        return self.current_tool
    
    def get_current_color(self):
        """获取当前选中的颜色"""
        return self.current_color
    
    def get_current_class(self):
        """获取当前选中的类别"""
        return self.current_class 